#include"compute.h"

At op(At x,Wt y){
    if(y==(Wt)1)
    	return x;
    else
    	return -x;
}

void conv2d(volatile At* in1,volatile At* in2,const Wt weight[Tm][Tn][K][K],Ot out[Tm][MAX_LEN*MAX_LEN],
		int fsize,int ch_in){
#pragma HLS array_partition variable=weight complete dim=0
#pragma HLS array_partition variable=out complete dim=1
	//linebuffer and window defintion
	hls::LineBuffer<K,MAX_LEN+2,At> LineBuffer1;
	hls::LineBuffer<K,MAX_LEN+2,At> LineBuffer2;
	hls::Window<K,K,At> WindowBuffer1;
	hls::Window<K,K,At> WindowBuffer2;
    //variable
    At tmp1;
    At tmp2;
    int addr=0;               //ofm address
    int start=(K-1)*fsize+(K-1);           //2=K-1,window is valid when n>=start
    //start compute
	for(int n=0;n<fsize*fsize;n++){
#pragma HLS LOOP_TRIPCOUNT min=14*14 max=14*14 avg=14*14
#pragma HLS PIPELINE II=1
#pragma HLS dependence variable=LineBuffer1 false
#pragma HLS dependence variable=WindowBuffer1 false
#pragma HLS dependence variable=LineBuffer2 false
#pragma HLS dependence variable=WindowBuffer2 false
            //deal with linebuffer
			tmp1=*(in1+ch_in*fsize*fsize+n);
			tmp2=*(in2+(ch_in+1)*fsize*fsize+n);             //����ͨ����ʼ��ַ
			int c=n%fsize;
    		LineBuffer1.shift_pixels_up(c);
    		LineBuffer2.shift_pixels_up(c);
            LineBuffer1.insert_bottom_row(tmp1,c);
            LineBuffer2.insert_bottom_row(tmp2,c);
            //deal with window
            WindowBuffer1.shift_pixels_left();
            WindowBuffer2.shift_pixels_left();
            At new_col1[K]={LineBuffer1.getval(0,c),LineBuffer1.getval(1,c),LineBuffer1.getval(2,c)};
            At new_col2[K]={LineBuffer2.getval(0,c),LineBuffer2.getval(1,c),LineBuffer2.getval(2,c)};
            WindowBuffer1.insert_right_col(new_col1);
            WindowBuffer2.insert_right_col(new_col2);
            //calc
            if(n>=start&&(n-start)%fsize<fsize-K+1){
            	Ot sum[Tm];
            	for(int m=0;m<Tm;m++){
            		sum[m]=(Ot)0;
            		for(int i=0;i<K;i++)
            			for(int j=0;j<K;j++)
            			{
            				sum[m]+=(op(WindowBuffer1.getval(i,j),weight[m][0][i][j])
            					 +op(WindowBuffer2.getval(i,j),weight[m][1][i][j]));
            			}
            		out[m][addr]=sum[m];
            	}
            	addr++;
            }

    	}
}

void add_partial_sum(Ot obuffer[Tm][MAX_LEN*MAX_LEN],Ot pbuffer[Tm][MAX_LEN*MAX_LEN],int ofsize){

	for(int i=0;i<ofsize*ofsize;i++){
#pragma HLS LOOP_TRIPCOUNT min=144 max=144 avg=144
#pragma HLS PIPELINE II=1
		for(int m=0;m<Tm;m++)
		{
#pragma HLS UNROLL
			obuffer[m][i]+=pbuffer[m][i];
		}
	}
}

void compute_one_block(volatile At* in1,volatile At* in2,const Wt weight[M][N][K][K],Ot obuffer[Tm][MAX_LEN*MAX_LEN],
		int fsize,int ch_in,int ch_out){
	//compute o[ch_out:ch_out+Tm][:][:]
	Ot pbuffer1[Tm][MAX_LEN*MAX_LEN];
	Ot pbuffer2[Tm][MAX_LEN*MAX_LEN];
	Wt w[Tm][Tn][K][K];
#pragma HLS array_partition variable=pbuffer1 complete dim=1
#pragma HLS array_partition variable=pbuffer2 complete dim=1
#pragma HLS array_partition variable=w complete dim=0
	//
	for(int i=0;i<(fsize-K+1)*(fsize-K+1);i++){
#pragma HLS LOOP_TRIPCOUNT min=144 max=144 avg=144
#pragma HLS PIPELINE II=1
		for(int m=0;m<Tm;m++){
#pragma HLS UNROLL
			obuffer[m][i]=(Ot)0;
		}
	}
	//
	load_weight(weight,w,0,ch_out);           //load weight[ch_out:ch_out+Tm][0:Tn][:][:]
	conv2d(in1,in2,w,pbuffer1,fsize,0);       //
	int pingpong=0;
	for(int ch=2;ch<ch_in;ch+=2){
#pragma HLS LOOP_TRIPCOUNT min=1 max=1 avg=1
		if(pingpong==0){
			load_weight(weight,w,ch,ch_out);
			conv2d(in1,in2,w,pbuffer2,fsize,ch);
			add_partial_sum(obuffer,pbuffer1,fsize-K+1);
			pingpong=1;
		}
		else{
			load_weight(weight,w,ch,ch_out);
			conv2d(in1,in2,w,pbuffer1,fsize,ch);
			add_partial_sum(obuffer,pbuffer2,fsize-K+1);
			pingpong=0;
		}
	}
	if(pingpong==0){
		add_partial_sum(obuffer,pbuffer1,fsize-K+1);
	}
	else{
		add_partial_sum(obuffer,pbuffer2,fsize-K+1);
	}
}
